from transformers import pipeline
from langchain_community.llms import HuggingFacePipeline
import time
hf = HuggingFacePipeline.from_model_id(
    model_id="THUDM/chatglm3-6b",
    task="text-generation",
    device=0,
    model_kwargs={"trust_remote_code":True},
    pipeline_kwargs={"max_new_tokens": 500},
)

for i in range(10):
    a = time.time()
    print(hf.invoke("西游记中描写了哪些人物"))
    print(time.time()-a)




